import json
import chardet
import matplotlib as mpl
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
from datetime import datetime, time
from matplotlib import ticker
from matplotlib.lines import Line2D
from matplotlib.patches import FancyArrowPatch
from adjustText import adjust_text

# 全局配置
mpl.rcParams.update({
    'font.family': 'Microsoft YaHei',
    'axes.unicode_minus': False,
    'agg.path.chunksize': 20000
})

# 文件路径
kline_trade_file = 'D:\\work\\code\\clark\\gitee\\big_a\\datas\\backtest\\sh513180\\20250714-202507161500_chart.txt'

def detect_encoding(file_path):
    with open(file_path, 'rb') as f:
        return chardet.detect(f.read(10000))['encoding']

def load_kline_data(file_path):
    kline_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"add":'):
                try:
                    kline_data.append(json.loads(line))
                except:
                    continue
    return kline_data

def load_trade_data(file_path):
    trade_data = []
    encoding = detect_encoding(file_path)
    with open(file_path, 'r', encoding=encoding, errors='replace') as file:
        for line in file:
            if line.strip().startswith('{"cash":'):
                try:
                    trade_data.append(json.loads(line))
                except:
                    continue
    return trade_data

# 加载数据
kline_data = load_kline_data(kline_trade_file)
trade_data = load_trade_data(kline_trade_file)

# 过滤交易时间段数据
def is_trading_time(dt):
    t = dt.time()
    return (time(9, 30) <= t <= time(11, 30)) or (time(13, 0) <= t <= time(15, 0))

# 准备K线数据
raw_dates = [datetime.strptime(k['date'], '%Y-%m-%d %H:%M:%S') for k in kline_data]
trading_mask = [is_trading_time(dt) for dt in raw_dates]

dates = [dt for dt, mask in zip(raw_dates, trading_mask) if mask]
opens = [k['open'] for k, mask in zip(kline_data, trading_mask) if mask]
highs = [k['high'] for k, mask in zip(kline_data, trading_mask) if mask]
lows = [k['low'] for k, mask in zip(kline_data, trading_mask) if mask]
closes = [k['close'] for k, mask in zip(kline_data, trading_mask) if mask]
volumes = [k['volume'] for k, mask in zip(kline_data, trading_mask) if mask]

# === 创建数值索引代替时间轴 ===
x_index = np.arange(len(dates))  # 连续的数值索引

# 创建图表
fig = plt.figure(figsize=(18, 14), layout="constrained")
gs = fig.add_gridspec(nrows=3, ncols=1, height_ratios=[4, 1, 1], hspace=0.05)

# 1. K线图
ax1 = fig.add_subplot(gs[0])
ax1.set_title(f'资产价格走势与交易点', fontsize=14, fontweight='bold')
ax1.set_ylabel('价格', fontsize=12)
ax1.grid(True, linestyle='--', alpha=0.7)

# 动态计算K线宽度
width = 0.6  # 固定宽度更稳定

# 绘制K线（使用数值索引）
for i in range(len(x_index)):
    color = 'red' if closes[i] >= opens[i] else 'green'
    ax1.vlines(x_index[i], lows[i], highs[i], color=color, linewidth=0.8)
    ax1.bar(x_index[i],
            height=abs(closes[i]-opens[i]),
            bottom=min(opens[i], closes[i]),
            width=width,
            color=color,
            edgecolor=color)

# 添加交易时段分割线（使用数值索引）
for i, dt in enumerate(dates):
    if dt.time() == time(13, 0):
        ax1.axvline(x_index[i], color='red', linestyle='--', alpha=0.4, linewidth=0.8)

# 2. 成交量（使用数值索引）
ax2 = fig.add_subplot(gs[1], sharex=ax1)
ax2.bar(x_index, volumes, width=width*0.8,
        color=np.where(np.array(closes) >= np.array(opens), 'red', 'green'))
ax2.set_ylabel('成交量', fontsize=12)
ax2.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{x/1e6:.1f}M'))
ax2.grid(True, linestyle='--', alpha=0.7)

# 3. 资产曲线
ax3 = fig.add_subplot(gs[2], sharex=ax1)
trade_dates = [datetime.strptime(t['date'], '%Y-%m-%d %H:%M:%S') for t in trade_data]
assets = [t['totalAsset'] for t in trade_data]

# 映射交易时间到K线索引位置
trade_indices = []
for td in trade_dates:
    min_idx = min(range(len(dates)), key=lambda i: abs(dates[i] - td))
    trade_indices.append(x_index[min_idx])

# 绘制资产曲线（使用数值索引）
ax3.plot(trade_indices, assets, 'b-', linewidth=2, alpha=0.8)
ax3.fill_between(trade_indices, min(assets), assets, color='blue', alpha=0.1)
ax3.set_ylabel('总资产', fontsize=12)
ax3.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'¥{x/1000:.0f}K'))
ax3.grid(True, linestyle='--', alpha=0.7)

# 标记起始点和终点（使用数值索引）
ax3.scatter(trade_indices[0], assets[0], s=80, color='green', edgecolor='white', zorder=5,
            label=f'起始: ¥{assets[0]:.3f}')
ax3.scatter(trade_indices[-1], assets[-1], s=80, color='red', edgecolor='white', zorder=5,
            label=f'结束: ¥{assets[-1]:.3f}')

# 标记资产峰值（使用数值索引）
peak_idx = np.argmax(assets)
ax3.scatter(trade_indices[peak_idx], assets[peak_idx], s=100, color='gold', edgecolor='darkorange',
            marker='*', zorder=6, label=f'峰值: ¥{assets[peak_idx]/1000:.0f}K')
ax3.legend(loc='upper left', fontsize=9)

# === 自定义刻度设置（显示关键时间点）===
tick_positions = []
tick_labels = []
current_day = dates[0].date() if dates else None
start_index = 0

for i, dt in enumerate(dates):
    if dt.date() != current_day:
        # 标记前一天结束点
        if i > 0:
            tick_positions.append(i-1)
            tick_labels.append(dates[i-1].strftime('%m-%d %H:%M'))
        # 标记新一天开始点
        tick_positions.append(i)
        tick_labels.append(dt.strftime('%m-%d %H:%M'))
        current_day = dt.date()
        start_index = i

# 标记最后一天结束点
if dates:
    tick_positions.append(len(dates)-1)
    tick_labels.append(dates[-1].strftime('%m-%d %H:%M'))

# 应用刻度设置
for ax in [ax1, ax2, ax3]:
    ax.set_xticks(tick_positions)
    ax.set_xticklabels(tick_labels, rotation=45, ha='right')

# 识别买卖点
buy_points, sell_points = [], []
for trade in trade_data:
    dt = datetime.strptime(trade['date'], '%Y-%m-%d %H:%M:%S')
    price = trade['price']
    operation = trade['operation']
    volume = trade.get('volume', trade.get('position', 0))

    min_idx = min(range(len(dates)), key=lambda i: abs(dates[i] - dt))
    if "买入" == operation:
        buy_points.append((x_index[min_idx], price, volume, min_idx))
    elif "卖出" == operation:
        sell_points.append((x_index[min_idx], price, volume, min_idx))

# 绘制买卖点标记（使用数值索引）
texts = []
for point in buy_points:
    x_pos, price, volume, idx = point
    t = ax1.text(x_pos, price, f'B {price:.3f}\n+{volume}',
                 fontsize=9, color='red', ha='center', va='top',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='red', alpha=0.8))
    texts.append(t)
    # 添加箭头
    ax1.add_patch(FancyArrowPatch(
        (x_pos, price*0.98), (x_pos, price),
        arrowstyle='->', color='red', lw=1, alpha=0.7))

for point in sell_points:
    x_pos, price, volume, idx = point
    t = ax1.text(x_pos, price, f'S {price:.3f}\n-{volume}',
                 fontsize=9, color='green', ha='center', va='bottom',
                 bbox=dict(boxstyle="round,pad=0.2", fc='white', ec='green', alpha=0.8))
    texts.append(t)
    # 添加箭头
    ax1.add_patch(FancyArrowPatch(
        (x_pos, price*1.02), (x_pos, price),
        arrowstyle='->', color='green', lw=1, alpha=0.7))

# 智能调整标注
adjust_text(texts,
            arrowprops=dict(arrowstyle='->', color='gray', lw=0.8, alpha=0.7, shrinkA=10),
            expand_points=(1.5, 1.5),
            expand_text=(1.2, 1.2),
            force_points=0.5,
            ax=ax1)

# 添加图例
legend_elements = [
    Line2D([0], [0], color='red', lw=4, label='阳线'),
    Line2D([0], [0], color='green', lw=4, label='阴线'),
    Line2D([0], [0], marker='o', color='red', label='买入点', markersize=8, linestyle='None'),
    Line2D([0], [0], marker='o', color='green', label='卖出点', markersize=8, linestyle='None'),
    Line2D([0], [0], color='red', linestyle='--', alpha=0.4, label='交易时段分割')
]
ax1.legend(handles=legend_elements, loc='upper left', fontsize=9)

# 调整布局
plt.subplots_adjust(left=0.06, right=0.95, top=0.95, bottom=0.12, hspace=0.15)
plt.savefig('专业K线分析图.png', dpi=150, bbox_inches='tight')
plt.show()